package edu.unika.aifb.graphindex.storage.lucene;

import java.io.File;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collection;
import java.util.Collections;
import java.util.HashSet;
import java.util.Iterator;
import java.util.List;
import java.util.Set;

import org.apache.log4j.Logger;
import org.apache.lucene.document.Document;
import org.apache.lucene.document.FieldSelector;
import org.apache.lucene.document.SetBasedFieldSelector;
import org.apache.lucene.index.CorruptIndexException;
import org.apache.lucene.index.IndexReader;
import org.apache.lucene.index.KeepOnlyLastCommitDeletionPolicy;
import org.apache.lucene.index.Term;
import org.apache.lucene.index.TermDocs;
import org.apache.lucene.index.TermEnum;
import org.apache.lucene.search.Collector;
import org.apache.lucene.search.IndexSearcher;
import org.apache.lucene.search.Query;
import org.apache.lucene.search.Scorer;
import org.apache.lucene.search.TermQuery;
import org.apache.lucene.store.FSDirectory;
import org.apache.lucene.store.NIOFSDirectory;

import edu.unika.aifb.graphindex.data.Table;
import edu.unika.aifb.graphindex.storage.DataField;
import edu.unika.aifb.graphindex.storage.IndexDescription;
import edu.unika.aifb.graphindex.storage.IndexStorageReader;
import edu.unika.aifb.graphindex.storage.StorageException;
import edu.unika.aifb.graphindex.util.StringSplitter;

public class LuceneIndexStorageReader implements IndexStorageReader {

	protected File m_dir;
	protected IndexDescription m_idx;
	protected IndexReader m_reader;
	protected IndexSearcher m_searcher;
	
	private static final Logger log = Logger.getLogger(IndexStorageReader.class);

	public LuceneIndexStorageReader(String directory, IndexDescription index) throws StorageException {
		this(new File(directory), index);
	}
	
	public LuceneIndexStorageReader(File dir, IndexDescription index) throws StorageException {
		m_dir = dir;
		m_idx = index;
		
		try {
			m_reader = IndexReader.open(NIOFSDirectory.open(m_dir), new KeepOnlyLastCommitDeletionPolicy(), true, LuceneUtil.LUCENE_TERMINFO_INDEX_DIVISOR);
			m_searcher = new IndexSearcher(m_reader);
		}
		catch (IOException e) {
			throw new StorageException(e);
		}
		fs = new SetBasedFieldSelector(new HashSet<String>(), new HashSet<String>(Arrays.asList(m_idx.getValueFieldName())));
	}
	
	public void close() throws StorageException {
		try {
			m_searcher.close();
			m_reader.close();
			m_reader = null;
			m_searcher = null;
		} catch (IOException e) {
			throw new StorageException(e);
		}
	}

	public void reopen() throws StorageException {
		try {
			m_searcher.close();
			m_reader.close();

			m_reader = IndexReader.open(FSDirectory.open(m_dir), true);
			m_searcher = new IndexSearcher(m_reader);
		} catch (IOException e) {
			throw new StorageException(e);
		}
	}
	
	protected IndexReader getReader() {
		return m_reader;
	}

	protected List<Integer> getDocumentIds(Query q) throws StorageException {
		final List<Integer> docIds = new ArrayList<Integer>();
		try {
			m_searcher.search(q, new Collector() {
				private int docBase;
				
				public void collect(int docId) {
					docIds.add(docId + docBase);
				}

				@Override
				public boolean acceptsDocsOutOfOrder() {
					return true;
				}

				@Override
				public void setNextReader(IndexReader reader, int docBase) throws IOException {
					this.docBase = docBase;
				}

				@Override
				public void setScorer(Scorer scorer) throws IOException {
				}
			});
		} catch (IOException e) {
			throw new StorageException(e);
		}
		
		Collections.sort(docIds);
		
		return docIds;
	}
	
	FieldSelector fs;
	protected Document getDocument(int docId) throws StorageException {
		try {
			Document doc = m_reader.document(docId, fs);
			return doc;
		} catch (CorruptIndexException e) {
			throw new StorageException(e);
		} catch (IOException e) {
			throw new StorageException(e);
		}
	}
	
	protected List<String> getDocumentValues(int docId) throws StorageException {
		Document doc = getDocument(docId);
		List<String> values;
		values = new ArrayList<String>();

		String valuesString = doc.getFieldable(m_idx.getValueFieldName()).stringValue();
		
		StringSplitter splitter = new StringSplitter(valuesString, "\n");
		String s;
		while ((s = splitter.next()) != null) 
			values.add(s);
	
		return values;
	}

	private void loadDocument(Table<String> table, int docId, IndexDescription index, int[] valueIdxs, String[] indexValues) throws StorageException {
		Document doc = getDocument(docId);
		
		String values = doc.getFieldable(index.getValueFieldName()).stringValue();
		
		// hack if the value field is empty, mainly for triple data (instead of quads)
		if (values.length() == 0) {
			throw new RuntimeException("wut");
		}
		
		StringSplitter splitter = new StringSplitter(values, "\n");
		
		String s;
		while ((s = splitter.next()) != null) {
			String[] row = new String [table.columnCount()];
			
			for (int i = 0; i < table.columnCount(); i++) {
				int idx = valueIdxs[i];
				
				if (idx == -1)
					row[i] = s;
				else
					row[i] = indexValues[idx];
			}
			table.addRow(row);
		}
	}

	private void getData(DataField field, Collection<String> values, String... indexFieldValues) throws StorageException {
		Table<String> table = getTable(new DataField[] { field }, indexFieldValues);
		for (String[] row : table)
			values.add(row[0]);
	}

	public List<String> getDataList(DataField field, String... indexFieldValues) throws StorageException {
		List<String> values = new ArrayList<String>();
		getData(field, values, indexFieldValues);
		return values;
	}

	public Set<String> getDataSet(DataField field, String... indexFieldValues) throws StorageException {
		Set<String> values = new HashSet<String>();
		getData(field, values, indexFieldValues);
		return values;
	}

	public String getDataItem(String... indexFieldValues) throws StorageException {
		if (indexFieldValues.length < m_idx.getIndexFields().size())
			throw new UnsupportedOperationException("getDataItem supports only queries with one result document");
		
		String value = null;
		
		TermQuery q = new TermQuery(new Term(m_idx.getIndexFieldName(), LuceneUtil.toIndexKey(indexFieldValues)));
		
		List<Integer> docIds = getDocumentIds(q);
		if (docIds.size() > 0) {
			Document doc = getDocument(docIds.get(0));
			value = doc.getFieldable(m_idx.getValueFieldName()).stringValue().trim();
		}

		return value;
	}
	
	public Table<String> getTable(DataField[] columns, String... indexFieldValues) throws StorageException {
		List<String> cols = new ArrayList<String>();
		for (DataField df : columns)
			cols.add(df.toString());
		Table<String> table = new Table<String>(cols);
		
		// build mapping of fields in index term to the desired output
		// columns in the table
		boolean usesValueField = false;
		int[] valueIdxs = new int [columns.length];
		for (int i = 0; i < columns.length; i++) {
			DataField colField = columns[i];
			
			if (colField == m_idx.getValueField()) {
				valueIdxs[i] = -1;
				usesValueField = true;
			}
			else {
				valueIdxs[i] = m_idx.getIndexFieldPos(colField);
			}
		}

		String indexKey = LuceneUtil.toIndexKey(indexFieldValues);

		try {
			TermEnum te = m_reader.terms(new Term(m_idx.getIndexFieldName(), indexKey));
			
			do {
				Term t = te.term();
				
				if (t == null || !t.text().startsWith(indexKey))
					break;
		
				String[] indexTerms = t.text().split(LuceneUtil.KEY_DELIM);

				if (!usesValueField) {
					// if the user does not need the value field, we can answer 
					// the query directly from the terms by splitting them
					
					String[] row = new String[table.columnCount()];
					for (int i = 0; i < row.length; i++)
						row[i] = indexTerms[valueIdxs[i]];
					
					table.addRow(row);
				}
				else {
					// otherwise, we retrieve all documents for the current term (for
					// a merged index there should be only one, but we want to deal
					// with unmerged indexes as well)
					TermDocs td = m_reader.termDocs(t);
					while (td.next()) 
						loadDocument(table, td.doc(), m_idx, valueIdxs, indexTerms);
				}
			}
			while (te.next());
		} catch (IOException e) {
			throw new StorageException(e);
		}
		return table;
	}

	public Table<String> getTable(DataField col1, DataField col2, String... indexFieldValues) throws StorageException {
		if (indexFieldValues.length < m_idx.getIndexFields().size())
			throw new UnsupportedOperationException("For this getTable method all index fields need to be specified (no prefixes).");
		
		if (m_idx.getValueField() != col1 && m_idx.getValueField() != col2)
			throw new UnsupportedOperationException("This getTable method requires one column to be the value field of the index.");
		
		Table<String> table = getTable(new DataField[] { col1, col2 }, indexFieldValues);
		
		if (m_idx.getValueField() == col1)
			table.setSortedColumn(0);
		else
			table.setSortedColumn(1);
			
		return table;
	}

	public class TermIterator implements Iterator<String[]> {
		private String m_indexKey;
		private TermEnum m_termEnum = null;
		private List<String> m_values;
		private Iterator<String> m_valueIterator;
		private int[] m_colIdx2TermIdx;
		private boolean m_usesValue = false;
		private String[] m_indexTerms;
		private Term m_next = null;
		private boolean m_empty = false;
		
		public TermIterator(DataField[] columnFields, String[] indexFieldValues) throws StorageException {
			m_indexKey = LuceneUtil.toIndexKey(indexFieldValues);
			
			try {
				m_termEnum = m_reader.terms(new Term(m_idx.getIndexFieldName(), m_indexKey));
				m_next = m_termEnum.term();
				
				if (!isPrefix(m_next))
					m_empty = true;
			} catch (IOException e) {
				throw new StorageException(e);
			}

			m_colIdx2TermIdx = new int [columnFields.length];
			for (int i = 0; i < columnFields.length; i++) {
				m_colIdx2TermIdx[i] = m_idx.getIndexFieldPos(columnFields[i]);
				if (columnFields[i] == m_idx.getValueField()) {
					m_usesValue = true;
					m_colIdx2TermIdx[i] = -1;
				}
			}
		}
		
		private boolean isPrefix(Term t) {
			if (t == null)
				return false;
			return t.text().startsWith(m_indexKey);
		}
		
		private Term peekNext() throws IOException {
			if (m_next == null) {
				m_termEnum.next();
				m_next = m_termEnum.term();
			}
			if (isPrefix(m_next))
				return m_next;
			else
				return null;
		}
		
		private Term getNext() throws IOException {
			if (m_next != null) {
				if (isPrefix(m_next)) {
					Term t = m_next;
					m_next = null;
					return t;
				}
				else
					return null;
			}
			else {
				if (!m_termEnum.next())
					return null;
				Term t = m_termEnum.term();
				if (isPrefix(t))
					return t;
				else
					return null;
			}
		}
		
		public boolean hasNext() {
			try {
				if (m_empty)
					return false;
				
				if (!m_usesValue && peekNext() == null)
					return false;

				if (m_usesValue && peekNext() == null && m_valueIterator != null && !m_valueIterator.hasNext())
					return false;
			} catch (IOException e) {
				e.printStackTrace();
				return false;
			}
			
			return true;
		}

		public String[] next() {
			if (m_indexTerms == null || m_values == null ||!m_valueIterator.hasNext()) {
				try {
					Term term = getNext();
					if (term == null)
						return null;
	
					String termText = term.text();
					m_indexTerms = termText.split(LuceneUtil.KEY_DELIM);
					
					if (m_indexTerms.length < m_idx.getIndexFields().size()) {
						String s = "";
						for (String idxTerm : m_indexTerms)
							s += idxTerm + " ";
						
						term = getNext();
						if (term == null)
							return null;
	
						termText = term.text();
						m_indexTerms = termText.split(LuceneUtil.KEY_DELIM);
					}
				}
				catch (IOException e) {
					e.printStackTrace();
					return null;
				}
			}
			
			String[] row = new String [m_colIdx2TermIdx.length];

			if (!m_usesValue) {
				for (int i = 0; i < m_colIdx2TermIdx.length; i++)
					row[i] = m_indexTerms[m_colIdx2TermIdx[i]];
			}
			else {
				if (m_values == null || !m_valueIterator.hasNext()) {
					try {
						m_values = getDataList(m_idx.getValueField(), m_indexTerms);
						m_valueIterator = m_values.iterator();
					} catch (StorageException e) {
						e.printStackTrace();
						return null;
					}
				}
				
				for (int i = 0; i < m_colIdx2TermIdx.length; i++) {
					if (m_colIdx2TermIdx[i] < 0)
						row[i] = m_valueIterator.next();
					else
						row[i] = m_indexTerms[m_colIdx2TermIdx[i]];
				}
			}
			
			return row;
		}

		public void remove() {
			throw new UnsupportedOperationException("remove not supported");
		}
	}


	public Iterator<String[]> iterator(DataField[] columns) throws StorageException {
		return new TermIterator(columns, new String[] {});
	}
	
	public Iterator<String[]> iterator(DataField[] columns, String... indexFieldValues) throws StorageException {
		return new TermIterator(columns, indexFieldValues);
	}
}
